import argparse
import warnings
warnings.filterwarnings('ignore')
import copy
import logging
import math
import os
import random
import shutil
from contextlib import nullcontext
from pathlib import Path

import accelerate
import numpy as np
import torch
import transformers
from accelerate import Accelerator
from accelerate.logging import get_logger
from accelerate.utils import DistributedType, ProjectConfiguration, set_seed

from datasets import load_dataset
from huggingface_hub import create_repo, upload_folder
from packaging import version
from peft import LoraConfig, set_peft_model_state_dict
from peft.utils import get_peft_model_state_dict
from PIL import Image
from torchvision import transforms
from tqdm.auto import tqdm

import diffusers
from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler, FluxFillPipeline, FluxTransformer2DModel
from diffusers.optimization import get_scheduler
from diffusers.training_utils import (
    cast_training_params,
    compute_density_for_timestep_sampling,
    compute_loss_weighting_for_sd3,
    free_memory,
)
from diffusers.utils import check_min_version, is_wandb_available, load_image, make_image_grid
from diffusers.utils.hub_utils import load_or_create_model_card, populate_model_card
from diffusers.utils.torch_utils import is_compiled_module
from diffusers.image_processor import VaeImageProcessor
from accelerate.utils import DistributedDataParallelKwargs

import pdb

from train_flux_parse import parse_args
from util_flux import pad_image
target_shape = (1024,1024)

# logger = logging.getLogger(__name__)
# logger.setLevel(logging.INFO)
# handler = logging.StreamHandler()
# logger.addHandler(handler)

# if is_wandb_available():
#     import wandb

# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
# check_min_version("0.33.0.dev0")

logger = get_logger(__name__)

NORM_LAYER_PREFIXES = ["norm_q", "norm_k", "norm_added_q", "norm_added_k"]

def pil2torch(img_path='',img_pil=None):
    if img_pil is None:
        img_pil = Image.open(img_path).convert("RGB")
    img_pil,_,_,_,_ = pad_image(img_pil)
    img_pil = img_pil.resize(target_shape)
    # pil_image = Image.open("input.jpg")
    image_np = np.array(img_pil)  # (H, W, 3)
    image_tensor = torch.from_numpy(image_np).permute(2, 0, 1).float() / 255.0  # (3, H, W)
    image_tensor = image_tensor.unsqueeze(0).to("cuda")  # (1, 3, H, W)
    return image_tensor
    
def encode_images(pixels: torch.Tensor, vae: torch.nn.Module, weight_dtype):
    pixel_latents = vae.encode(pixels.to(vae.dtype)).latent_dist.sample()
    pixel_latents = (pixel_latents - vae.config.shift_factor) * vae.config.scaling_factor
    return pixel_latents.to(weight_dtype)


def get_train_dataset(args, accelerator):
    dataset = None
    if args.dataset_name is not None:
        dataset = load_dataset(
            args.dataset_name,
            data_files=args.dataset_config_name,
        )
    if args.jsonl_for_train is not None:
        # load from json
        dataset = load_dataset("json", 
                               data_files=args.jsonl_for_train, 
                               field = 'data',
                               cache_dir=args.cache_dir)
        dataset = dataset.flatten_indices()
    # Preprocessing the datasets.
    # We need to tokenize inputs and targets.
    column_names = dataset["train"].column_names


    # 6. Get the column names for input/target.
    if args.image_column is None:
        image_column = column_names[0]
        logger.info(f"image column defaulting to {image_column}")
    else:
        image_column = args.image_column
        # if image_column not in column_names:
        #     raise ValueError(
        #         f"`--image_column` value '{args.image_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}"
        #     )

    if args.deep_image_column is None:
        deep_image_column = column_names[1]
        logger.info(f"deep image column defaulting to {deep_image_column}")
    else:
        deep_image_column = args.deep_image_column
        # if deep_image_column not in column_names:
        #     raise ValueError(
        #         f"`--deep_image_column` value '{args.deep_image_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}"
        #     )
        
    if args.texture_image_column is None:
        texture_image_column = column_names[2]
        logger.info(f"texture image column defaulting to {texture_image_column}")
    else:
        texture_image_column = args.texture_image_column
        # if deep_image_column not in column_names:
        #     raise ValueError(
        #         f"`--texture_image_column` value '{args.texture_image_column}' not found in dataset columns. Dataset columns are: {', '.join(column_names)}"
        #     )

    with accelerator.main_process_first():
        train_dataset = dataset["train"].shuffle(seed=args.seed)
        if args.max_train_samples is not None:
            train_dataset = train_dataset.select(range(args.max_train_samples))
    return train_dataset


def prepare_train_dataset(dataset, accelerator,image_processor, mask_processor):
    def cat_image(ref_image, image):
        ref_image = ref_image.resize(image.size)
        output_image = Image.new('RGB', (image.size[0]*2, 
                                         image.size[1]), (0,0,0))
        output_image.paste(ref_image, (0, 0))
        output_image.paste(image, (image.size[0], 0))
        return output_image
    
    
    from util_flux import process_img_1024
    def horizontal_concat_images(images,):
            # 计算总宽度 (所有图片宽度之和 + 间距)
            widths, heights = zip(*(img.size for img in images))
            total_width = sum(widths)
            max_height = max(heights)
            
            # 创建新图片 (白色背景)
            new_img = Image.new('RGB', (total_width, max_height), color='white')
            
            # 粘贴图片并绘制分隔线
            x_offset = 0
            for i, img in enumerate(images):
                new_img.paste(img, (x_offset, 0))
                
                x_offset += img.width
            
            return new_img
    def preprocess_train(examples):
        # cloth
        images = [
            (process_img_1024(image))
            for image in examples['clo_path']
        ]
        # human
        human_images = [
            (process_img_1024(image))
            for image in examples['human_path']
        ]

        images = [ horizontal_concat_images([i1,i2])
            for i1,i2 in zip(images,human_images)]

        # numpy
        images = [image_processor.preprocess(image, height=args.resolution_height, width=args.resolution_width) 
                  for image in images]
        

        black_image = Image.new("RGB", (args.resolution_width, args.resolution_height), (0, 0, 0))
        white_image = Image.new("RGB", (args.resolution_width//2, args.resolution_height), (255, 255, 255))
        # black_image.paste(white_image, (args.resolution_width//2, 0))
        black_image.paste(white_image, (0, 0))
        
        mask_images = [black_image]
        mask_images = [mask_processor.preprocess(black_image,
                                                 height=args.resolution_height, 
                                                 width=args.resolution_width)]
        
        # 保留 mask 中 黑色部分的condition，白色部分的待生成部分mask掉
        masked_images = [images[idex] * (1 - mask_images[idex]) 
                         for idex in range(len(mask_images))] 

        # mask_images = [
        #     black_image
        #     for image in examples[args.ref_image_column]
        # ]

        # mask_images = [mask_processor.preprocess(image, 
        #                                          height=args.resolution_height, 
        #                                          width=args.resolution_width) 
        #                                          for image in mask_images]

        # masked_images = [texture_images[idex] * (1 - mask_images[idex]) 
        #                  for idex in range(len(mask_images))]

        examples["pixel_values"] = images
        # examples["deep_pixel_values"] = deep_images
        examples["mask_pixel_values"] = mask_images # 左黑右白mask
        examples["masked_pixel_values"] = masked_images # mask*image => mask*texture
        ## extra texture!
        # examples["texture_pixel_values"] = texture_images
        # emb_path = examples['concat_emb']
        # pdb.set_trace()
        # print(emb_path)
        # emb = torch.load(emb_path,map_location=accelerator.device)
        
        # examples['prompt_embeds'] = emb['prompt_embeds'].to(dtype = accelerator.dtype)
        # examples['pooled_prompt_embeds'] = emb['pooled_prompt_embeds'].to(dtype = accelerator.dtype)
        # examples['text_ids'] = emb['text_ids'].to(dtype = accelerator.dtype)

        return examples

    with accelerator.main_process_first():
        dataset = dataset.with_transform(preprocess_train)

    return dataset


def collate_fn(examples):
    pixel_values = torch.stack([example["pixel_values"] for example in examples])
    pixel_values = pixel_values.to(memory_format=torch.contiguous_format).float()
    mask_pixel_values = torch.stack([example["mask_pixel_values"] for example in examples])
    mask_pixel_values = mask_pixel_values.to(memory_format=torch.contiguous_format).float()
    masked_pixel_values = torch.stack([example["masked_pixel_values"] for example in examples])
    masked_pixel_values = masked_pixel_values.to(memory_format=torch.contiguous_format).float()

    
    return {"pixel_values": pixel_values, 
            "mask_pixel_values": mask_pixel_values, 
            "masked_pixel_values": masked_pixel_values, 
            "emb_path": examples[0]['emb_path'],


            }


def main(args):
    if args.use_lora_bias and args.gaussian_init_lora:
        raise ValueError("`gaussian` LoRA init scheme isn't supported when `use_lora_bias` is True.")

    logging_out_dir = Path(args.output_dir, args.logging_dir)

    if torch.backends.mps.is_available() and args.mixed_precision == "bf16":
        # due to pytorch#99272, MPS does not yet support bfloat16.
        raise ValueError(
            "Mixed precision training with bfloat16 is not supported on MPS. Please use fp16 (recommended) or fp32 instead."
        )

    accelerator_project_config = ProjectConfiguration(project_dir=args.output_dir, 
                                                      logging_dir=str(logging_out_dir))

    accelerator = Accelerator(
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        mixed_precision=args.mixed_precision,
        log_with=args.report_to,
        project_config=accelerator_project_config,
        # debug=True
        # kwargs_handlers=[DistributedDataParallelKwargs(find_unused_parameters=True)]
    )

    # Disable AMP for MPS. A technique for accelerating machine learning computations on iOS and macOS devices.
    if torch.backends.mps.is_available():
        logger.info("MPS is enabled. Disabling AMP.")
        accelerator.native_amp = False

    # Make one log on every process with the configuration for debugging.
    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        # DEBUG, INFO, WARNING, ERROR, CRITICAL
        level=logging.INFO,
    )
    logger.info(accelerator.state, main_process_only=False)

    if accelerator.is_local_main_process:
        transformers.utils.logging.set_verbosity_warning()
        diffusers.utils.logging.set_verbosity_info()
    else:
        transformers.utils.logging.set_verbosity_error()
        diffusers.utils.logging.set_verbosity_error()

    # If passed along, set the training seed now.
    if args.seed is not None:
        set_seed(args.seed)

    # Handle the repository creation
    if accelerator.is_main_process:
        if args.output_dir is not None:
            os.makedirs(args.output_dir, exist_ok=True)

        if args.push_to_hub:
            repo_id = create_repo(
                repo_id=args.hub_model_id or Path(args.output_dir).name, 
                exist_ok=True, token=args.hub_token
            ).repo_id

    # cast down and move to the CPU
    weight_dtype = torch.float32
    if accelerator.mixed_precision == "fp16":
        weight_dtype = torch.float16
    elif accelerator.mixed_precision == "bf16":
        weight_dtype = torch.bfloat16
    

    # Load models. We will load the text encoders later in a pipeline to compute
    # embeddings.
    vae = AutoencoderKL.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="vae",
        revision=args.revision,
        variant=args.variant,
    )
    vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
    image_processor = VaeImageProcessor( # vae channel 4 scale 16
        vae_scale_factor=vae_scale_factor * 2)
    latent_channels = vae.config.latent_channels
    mask_processor = VaeImageProcessor( # vae channel 16 scale 16
            vae_scale_factor=vae_scale_factor * 2,
            vae_latent_channels=latent_channels,
            do_normalize=False,
            do_binarize=True,
            do_convert_grayscale=True,
        )
    flux_transformer = FluxTransformer2DModel.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="transformer",
        revision=args.revision,
        variant=args.variant,
    )
    logger.info("All models loaded successfully")

    print("loaded pipeline")

    noise_scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="scheduler",
    )
    noise_scheduler_copy = copy.deepcopy(noise_scheduler)
    vae.requires_grad_(False)
    flux_transformer.requires_grad_(False)

    # let's not move the VAE to the GPU yet.
    vae.to(device=accelerator.device,dtype=torch.float32)  # keep the VAE in float32.

    # enable image inputs
    with torch.no_grad():
        initial_input_channels = flux_transformer.config.in_channels
        initial_output_channels = flux_transformer.config.out_channels
        new_linear = torch.nn.Linear( # ori_in 64 now_in 74 out 3072
            # flux_transformer.x_embedder.in_features * 3,
            flux_transformer.x_embedder.in_features,
            flux_transformer.x_embedder.out_features,
            bias=flux_transformer.x_embedder.bias is not None,
            dtype=flux_transformer.dtype,
            device=flux_transformer.device,
        )


        new_linear.weight.zero_()
        new_linear.weight[:, :initial_input_channels].copy_(flux_transformer.x_embedder.weight)
        if flux_transformer.x_embedder.bias is not None:
            new_linear.bias.copy_(flux_transformer.x_embedder.bias)
        flux_transformer.x_embedder = new_linear

    assert torch.all(flux_transformer.x_embedder.weight[:, initial_input_channels:].data == 0)
    flux_transformer.register_to_config(
        # in_channels=initial_input_channels * 3, 
        in_channels=initial_input_channels, 
        out_channels=initial_output_channels)

    if args.quant_transformers:
        # from torchao.quantization import quantize_, int8_weight_only
        # quantize_(flux_transformer, int8_weight_only())

        from optimum.quanto import freeze, qfloat8, quantize, QTensor, qint4
        quantize(flux_transformer, weights=qfloat8)
        freeze(flux_transformer)
        logger.info("quant transformer")
        # flux_transformer.to(accelerator.device)
    else:
        flux_transformer.to(accelerator.device, dtype=weight_dtype)

    if args.train_norm_layers:
        for name, param in flux_transformer.named_parameters():
            if any(k in name for k in NORM_LAYER_PREFIXES):
                param.requires_grad = True

    if args.lora_layers is not None:
        if args.lora_layers != "all-linear":
            target_modules = [layer.strip() for layer in args.lora_layers.split(",")]
            # add the input layer to the mix.
            if "x_embedder" not in target_modules:
                target_modules.append("x_embedder")
        elif args.lora_layers == "all-linear":
            target_modules = set()
            for name, module in flux_transformer.named_modules():
                if isinstance(module, torch.nn.Linear):
                    target_modules.add(name)
            target_modules = list(target_modules)
    else:
        target_modules = [
            "x_embedder",
            "attn.to_k",
            "attn.to_q",
            "attn.to_v",
            "attn.to_out.0",
            "attn.add_k_proj",
            "attn.add_q_proj",
            "attn.add_v_proj",
            "attn.to_add_out",
            "ff.net.0.proj",
            "ff.net.2",
            "ff_context.net.0.proj",
            "ff_context.net.2",
        ]
    transformer_lora_config = LoraConfig(
        r=args.rank,
        lora_alpha=args.rank,
        init_lora_weights="gaussian" if args.gaussian_init_lora else True,
        target_modules=target_modules,
        lora_bias=args.use_lora_bias,
    )
    flux_transformer.add_adapter(transformer_lora_config)

    if False:
        from safetensors.torch import load_file 
        lora_load_path = "/data/models/train/fabric/fabric_coat/pytorch_lora_weights.safetensors"
        loaded_weights = load_file(lora_load_path, device="cpu")
        original_grad_state = {name: param for name, param in flux_transformer.named_parameters()}
        for name, param in original_grad_state.items():
            if param.requires_grad:
                new_name = name.replace("default.", "")
                new_name = 'transformer.' + new_name
                param.data.copy_(loaded_weights[new_name])

        flux_transformer.to(accelerator.device)

    def prepare_mask_latents(
        mask,
        masked_image,
        batch_size,
        num_channels_latents,
        num_images_per_prompt,
        height,
        width,
        dtype,
        device,
    ):
        # 1. calculate the height and width of the latents
        # VAE applies 8x compression on images but we must also account for packing which requires
        # latent height and width to be divisible by 2.
        height = 2 * (int(height) // (vae_scale_factor * 2))
        width = 2 * (int(width) // (vae_scale_factor * 2))

        # 2. encode the masked image
        if masked_image.shape[1] == num_channels_latents:
            masked_image_latents = masked_image
        else:
            masked_image_latents = vae.encode(masked_image.to(device=vae.device, dtype=vae.dtype)).latent_dist.sample()

        masked_image_latents = (masked_image_latents - vae.config.shift_factor) * vae.config.scaling_factor
        masked_image_latents = masked_image_latents.to(device=device, dtype=dtype)

        # 3. duplicate mask and masked_image_latents for each generation per prompt, using mps friendly method
        batch_size = batch_size * num_images_per_prompt
        if mask.shape[0] < batch_size:
            if not batch_size % mask.shape[0] == 0:
                raise ValueError(
                    "The passed mask and the required batch size don't match. Masks are supposed to be duplicated to"
                    f" a total batch size of {batch_size}, but {mask.shape[0]} masks were passed. Make sure the number"
                    " of masks that you pass is divisible by the total requested batch size."
                )
            mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1)
        if masked_image_latents.shape[0] < batch_size:
            if not batch_size % masked_image_latents.shape[0] == 0:
                raise ValueError(
                    "The passed images and the required batch size don't match. Images are supposed to be duplicated"
                    f" to a total batch size of {batch_size}, but {masked_image_latents.shape[0]} images were passed."
                    " Make sure the number of images that you pass is divisible by the total requested batch size."
                )
            masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1)

        # 5.resize mask to latents shape we we concatenate the mask to the latents
        mask = mask[:, 0, :, :]  # batch_size, 8 * height, 8 * width (mask has not been 8x compressed)
        mask = mask.view(
            batch_size, height, vae_scale_factor, width, vae_scale_factor
        )  # batch_size, height, 8, width, 8
        mask = mask.permute(0, 2, 4, 1, 3)  # batch_size, 8, 8, height, width
        mask = mask.reshape(
            batch_size, vae_scale_factor * vae_scale_factor, height, width
        )  # batch_size, 8*8, height, width
        mask = mask.to(device=device, dtype=dtype)

        return mask, masked_image_latents
    
    def unwrap_model(model):
        model = accelerator.unwrap_model(model)
        model = model._orig_mod if is_compiled_module(model) else model
        return model

    # `accelerate` 0.16.0 will have better support for customized saving
    if version.parse(accelerate.__version__) >= version.parse("0.16.0"):

        def save_model_hook(models, weights, output_dir):
            if accelerator.is_main_process:
                transformer_lora_layers_to_save = None

                for model in models:
                    if isinstance(unwrap_model(model), type(unwrap_model(flux_transformer))):
                        model = unwrap_model(model)
                        transformer_lora_layers_to_save = get_peft_model_state_dict(model)
                        if args.train_norm_layers:
                            transformer_norm_layers_to_save = {
                                f"transformer.{name}": param
                                for name, param in model.named_parameters()
                                if any(k in name for k in NORM_LAYER_PREFIXES)
                            }
                            transformer_lora_layers_to_save = {
                                **transformer_lora_layers_to_save,
                                **transformer_norm_layers_to_save,
                            }
                    else:
                        raise ValueError(f"unexpected save model: {model.__class__}")

                    # make sure to pop weight so that corresponding model is not saved again
                    if weights:
                        weights.pop()

                FluxFillPipeline.save_lora_weights(
                    output_dir,
                    transformer_lora_layers=transformer_lora_layers_to_save,
                )

        def load_model_hook(models, input_dir):
            transformer_ = None

            if not accelerator.distributed_type == DistributedType.DEEPSPEED:
                while len(models) > 0:
                    model = models.pop()

                    if isinstance(model, type(unwrap_model(flux_transformer))):
                        transformer_ = model
                    else:
                        raise ValueError(f"unexpected save model: {model.__class__}")
            else:
                transformer_ = FluxTransformer2DModel.from_pretrained(
                    args.pretrained_model_name_or_path, subfolder="transformer"
                ).to(accelerator.device, weight_dtype)

                # Handle input dimension doubling before adding adapter
                with torch.no_grad():
                    initial_input_channels = transformer_.config.in_channels
                    new_linear = torch.nn.Linear(
                        transformer_.x_embedder.in_features * 7 // 6,
                        transformer_.x_embedder.out_features,
                        bias=transformer_.x_embedder.bias is not None,
                        dtype=transformer_.dtype,
                        device=transformer_.device,
                    )
                    new_linear.weight.zero_()
                    new_linear.weight[:, :initial_input_channels].copy_(transformer_.x_embedder.weight)
                    if transformer_.x_embedder.bias is not None:
                        new_linear.bias.copy_(transformer_.x_embedder.bias)
                    transformer_.x_embedder = new_linear
                    transformer_.register_to_config(in_channels=initial_input_channels * 7 // 6)

                transformer_.add_adapter(transformer_lora_config)

            lora_state_dict = FluxFillPipeline.lora_state_dict(input_dir)
            transformer_lora_state_dict = {
                f'{k.replace("transformer.", "")}': v
                for k, v in lora_state_dict.items()
                if k.startswith("transformer.") and "lora" in k
            }
            incompatible_keys = set_peft_model_state_dict(
                transformer_, transformer_lora_state_dict, adapter_name="default"
            )
            if incompatible_keys is not None:
                # check only for unexpected keys
                unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None)
                if unexpected_keys:
                    logger.warning(
                        f"Loading adapter weights from state_dict led to unexpected keys not found in the model: "
                        f" {unexpected_keys}. "
                    )
            if args.train_norm_layers:
                transformer_norm_state_dict = {
                    k: v
                    for k, v in lora_state_dict.items()
                    if k.startswith("transformer.") and any(norm_k in k for norm_k in NORM_LAYER_PREFIXES)
                }
                transformer_._transformer_norm_layers = FluxFillPipeline._load_norm_into_transformer(
                    transformer_norm_state_dict,
                    transformer=transformer_,
                    discard_original_layers=False,
                )

            # Make sure the trainable params are in float32. This is again needed since the base models
            # are in `weight_dtype`. More details:
            # https://github.com/huggingface/diffusers/pull/6514#discussion_r1449796804
            if args.mixed_precision == "fp16":
                models = [transformer_]
                # only upcast trainable parameters (LoRA) into fp32
                cast_training_params(models)

        accelerator.register_save_state_pre_hook(save_model_hook)
        accelerator.register_load_state_pre_hook(load_model_hook)

    # Make sure the trainable params are in float32.
    if args.mixed_precision == "fp16":
        models = [flux_transformer]
        # only upcast trainable parameters (LoRA) into fp32
        cast_training_params(models, dtype=torch.float32)

    if args.gradient_checkpointing:
        flux_transformer.enable_gradient_checkpointing()

    # Enable TF32 for faster training on Ampere GPUs,
    # cf https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices
    if args.allow_tf32:
        torch.backends.cuda.matmul.allow_tf32 = True

    if args.scale_lr:
        args.learning_rate = (
            args.learning_rate * args.gradient_accumulation_steps * args.train_batch_size * accelerator.num_processes
        )

    # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs
    if args.use_8bit_adam:
        try:
            import bitsandbytes as bnb
        except ImportError:
            raise ImportError(
                "To use 8-bit Adam, please install the bitsandbytes library: `pip install bitsandbytes`."
            )

        optimizer_class = bnb.optim.AdamW8bit
    else:
        optimizer_class = torch.optim.AdamW

    # Optimization parameters
    transformer_lora_parameters = list(filter(lambda p: p.requires_grad, flux_transformer.parameters()))
    # print(transformer_lora_parameters)
    # import pdb
    optimizer = optimizer_class(
        transformer_lora_parameters,
        lr=args.learning_rate,
        betas=(args.adam_beta1, args.adam_beta2),
        weight_decay=args.adam_weight_decay,
        eps=args.adam_epsilon,
    )

    # Prepare dataset and dataloader.
    train_dataset = get_train_dataset(args, accelerator)
    train_dataset = prepare_train_dataset(train_dataset, accelerator, image_processor, mask_processor)
    train_dataloader = torch.utils.data.DataLoader(
        train_dataset,
        shuffle=True,
        collate_fn=collate_fn,
        batch_size=args.train_batch_size,
        num_workers=args.dataloader_num_workers,
    )

    # Scheduler and math around the number of training steps.
    # Check the PR https://github.com/huggingface/diffusers/pull/8312 for detailed explanation.
    if args.max_train_steps is None:
        len_train_dataloader_after_sharding = math.ceil(len(train_dataloader) / accelerator.num_processes)
        num_update_steps_per_epoch = math.ceil(len_train_dataloader_after_sharding / args.gradient_accumulation_steps)
        num_training_steps_for_scheduler = (
            args.num_train_epochs * num_update_steps_per_epoch * accelerator.num_processes
        )
    else:
        num_training_steps_for_scheduler = args.max_train_steps * accelerator.num_processes

    lr_scheduler = get_scheduler(
        args.lr_scheduler,
        optimizer=optimizer,
        num_warmup_steps=args.lr_warmup_steps * accelerator.num_processes,
        num_training_steps=args.max_train_steps * accelerator.num_processes,
        num_cycles=args.lr_num_cycles,
        power=args.lr_power,
    )
    # Prepare everything with our `accelerator`.
    flux_transformer, optimizer, train_dataloader, lr_scheduler = accelerator.prepare(
        flux_transformer, optimizer, train_dataloader, lr_scheduler
    )

    # We need to recalculate our total training steps as the size of the training dataloader may have changed.
    num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
    if args.max_train_steps is None:
        args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
        if num_training_steps_for_scheduler != args.max_train_steps * accelerator.num_processes:
            logger.warning(
                f"The length of the 'train_dataloader' after 'accelerator.prepare' ({len(train_dataloader)}) does not match "
                f"the expected length ({len_train_dataloader_after_sharding}) when the learning rate scheduler was created. "
                f"This inconsistency may result in the learning rate scheduler not functioning properly."
            )
    # Afterwards we recalculate our number of training epochs
    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)

    # We need to initialize the trackers we use, and also store our configuration.
    # The trackers initializes automatically on the main process.
    if accelerator.is_main_process:
        tracker_config = dict(vars(args))
        accelerator.init_trackers(args.tracker_project_name, config=tracker_config)

    # Train!
    total_batch_size = args.train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps

    logger.info("***** Running training *****")
    logger.info(f"  Num examples = {len(train_dataset)}")
    logger.info(f"  Num batches each epoch = {len(train_dataloader)}")
    logger.info(f"  Num Epochs = {args.num_train_epochs}")
    logger.info(f"  Instantaneous batch size per device = {args.train_batch_size}")
    logger.info(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    logger.info(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
    logger.info(f"  Total optimization steps = {args.max_train_steps}")
    global_step = 0
    first_epoch = 0

    # Potentially load in the weights and states from a previous save
    if args.resume_from_checkpoint:
        if args.resume_from_checkpoint != "latest":
            path = os.path.basename(args.resume_from_checkpoint)
        else:
            # Get the most recent checkpoint
            dirs = os.listdir(args.output_dir)
            dirs = [d for d in dirs if d.startswith("checkpoint")]
            dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
            path = dirs[-1] if len(dirs) > 0 else None

        if path is None:
            logger.info(f"Checkpoint '{args.resume_from_checkpoint}' does not exist. Starting a new training run.")
            args.resume_from_checkpoint = None
            initial_global_step = 0
        else:
            logger.info(f"Resuming from checkpoint {path}")
            accelerator.load_state(os.path.join(args.output_dir, path))
            global_step = int(path.split("-")[1])

            initial_global_step = global_step
            first_epoch = global_step // num_update_steps_per_epoch
    else:
        initial_global_step = 0

    progress_bar = tqdm(
        range(0, args.max_train_steps),
        initial=initial_global_step,
        desc="Steps",
        # Only show the progress bar once on each machine.
        disable=not accelerator.is_local_main_process,
    )

    def get_sigmas(timesteps, n_dim=4, dtype=torch.float32):
        sigmas = noise_scheduler_copy.sigmas.to(device=accelerator.device, dtype=dtype)
        schedule_timesteps = noise_scheduler_copy.timesteps.to(accelerator.device)
        timesteps = timesteps.to(accelerator.device)
        step_indices = [(schedule_timesteps == t).nonzero().item() for t in timesteps]

        sigma = sigmas[step_indices].flatten()
        while len(sigma.shape) < n_dim:
            sigma = sigma.unsqueeze(-1)
        return sigma

    for epoch in range(first_epoch, args.num_train_epochs):
        flux_transformer.train()
        for step, batch in enumerate(train_dataloader):
            torch.cuda.empty_cache()
            with accelerator.accumulate(flux_transformer):

                ## load emb
                # pdb.set_trace()
                emb_path = batch['emb_path']
                # concat_emb_path = batch['concat_emb']  # list
                # pdb.set_trace()
                concat_emb = torch.load(emb_path)
                prompt_embeds = concat_emb['prompt_embeds']\
                    .to(accelerator.device,dtype = weight_dtype)   
                pooled_prompt_embeds = concat_emb['pooled_prompt_embeds']\
                    .to(accelerator.device,dtype = weight_dtype)
                text_ids = concat_emb['text_ids']\
                    .to(accelerator.device,dtype = weight_dtype)
                # pdb.set_trace()


                # Convert images to latent space
                # vae encode
                pixel_latents = encode_images(batch["pixel_values"].squeeze(1), 
                                              vae, weight_dtype)
                # deep_latents = encode_images(
                #     batch["deep_pixel_values"].squeeze(1), 
                #     vae, weight_dtype
                # )
                # texture_latents = encode_images(
                #     batch["texture_pixel_values"].squeeze(1), 
                #     vae, weight_dtype
                # )
                masked_images = batch["masked_pixel_values"].squeeze(1).to(device=accelerator.device, 
                                                                           dtype=weight_dtype)
                mask_images = batch["mask_pixel_values"].squeeze(1).to(device=accelerator.device, 
                                                                       dtype=weight_dtype)
                
                # if args.offload:
                    # offload vae to CPU.
                    # vae.cpu()

                # Sample a random timestep for each image
                # for weighting schemes where we sample timesteps non-uniformly
                bsz = pixel_latents.shape[0]
                num_channels_latents = vae.config.latent_channels
                num_images_per_prompt = 1
                # torch.Size([1, 16, 128, 384])
                mask, masked_image_latents = prepare_mask_latents(
                    mask_images,
                    masked_images,
                    bsz,
                    num_channels_latents,
                    num_images_per_prompt,
                    args.resolution_height,
                    args.resolution_width,
                    weight_dtype,
                    accelerator.device,
                )
                
                # noise : torch.Size([1, 16, 128, 384]) 128*384/4=12288
                noise = torch.randn_like(pixel_latents, device=accelerator.device, dtype=weight_dtype)
                u = compute_density_for_timestep_sampling(
                    weighting_scheme=args.weighting_scheme,
                    batch_size=bsz,
                    logit_mean=args.logit_mean,
                    logit_std=args.logit_std,
                    mode_scale=args.mode_scale,
                )
                indices = (u * noise_scheduler_copy.config.num_train_timesteps).long()
                timesteps = noise_scheduler_copy.timesteps[indices].to(device=pixel_latents.device)

                # Add noise according to flow matching.
                sigmas = get_sigmas(timesteps, n_dim=pixel_latents.ndim, dtype=pixel_latents.dtype)
                noisy_model_input = (1.0 - sigmas) * pixel_latents + sigmas * noise
                # # Concatenate across channels.
                # # Question: Should we concatenate before adding noise?
                # concatenated_noisy_model_input = torch.cat([noisy_model_input, masked_image_latents, mask, deep_latents], dim=1)

                # pack the latents.
                packed_noisy_model_input = FluxFillPipeline._pack_latents(
                    noisy_model_input, # torch.Size([1, 12288, 64])
                    batch_size=bsz,
                    num_channels_latents=noisy_model_input.shape[1],
                    height=noisy_model_input.shape[2],
                    width=noisy_model_input.shape[3],
                )
                packed_masked_image_latents = FluxFillPipeline._pack_latents(
                    masked_image_latents, # torch.Size([1, 12288, 64])
                    batch_size=bsz,
                    num_channels_latents=masked_image_latents.shape[1],
                    height=masked_image_latents.shape[2],
                    width=masked_image_latents.shape[3],
                )
                packed_mask_image_latents = FluxFillPipeline._pack_latents(
                    mask,  # torch.Size([1, 12288, 256])
                    batch_size=bsz,
                    num_channels_latents=mask.shape[1],
                    height=mask.shape[2],
                    width=mask.shape[3],
                )
                # pdb.set_trace()
                # del packed_masked_image_latents,packed_mask_image_latents
                # packed_deep_latents = FluxFillPipeline._pack_latents(
                #     deep_latents,
                #     batch_size=bsz,
                #     num_channels_latents=deep_latents.shape[1],
                #     height=deep_latents.shape[2],
                #     width=deep_latents.shape[3],
                # )
                # packed_texture_latents = FluxFillPipeline._pack_latents(
                #     texture_latents,
                #     batch_size=bsz,
                #     num_channels_latents=texture_latents.shape[1],
                #     height=texture_latents.shape[2],
                #     width=texture_latents.shape[3],
                # )
                '''
                torch.Size([1, 16, 128, 128]) torch.Size([1, 4096, 64]) 16*4=64
                torch.Size([1, 16, 128, 128]) torch.Size([1, 4096, 64])
                torch.Size([1, 16, 128, 128]) torch.Size([1, 4096, 64])
                torch.Size([1, 16, 128, 128]) torch.Size([1, 4096, 64])
                torch.Size([1, 64, 128, 128]) torch.Size([1, 4096, 256])
                
                default-FLUX-Fill : torch.Size([1, 12288, 384])
                '''
                packed_noisy_model_input = torch.cat((packed_noisy_model_input, 
                                                    #   packed_deep_latents,
                                                    #   packed_texture_latents,
                                                      packed_masked_image_latents, 
                                                      packed_mask_image_latents,
                                                      ), dim=-1)
                

                # latent image ids for RoPE. torch.Size([4096, 3])
                latent_image_ids = FluxFillPipeline._prepare_latent_image_ids(
                    bsz,
                    noisy_model_input.shape[2] // 2,
                    noisy_model_input.shape[3] // 2,
                    accelerator.device,
                    weight_dtype,
                )

                # handle guidance
                if unwrap_model(flux_transformer).config.guidance_embeds:
                    guidance_vec = torch.full(
                        (bsz,),
                        args.guidance_scale,
                        device=noisy_model_input.device,
                        dtype=weight_dtype,
                    )
                else:
                    guidance_vec = None

                # Predict.
                # packed_noisy_model_input torch.Size([1, 4096, 192])
                # model_pred torch.Size([1, 256, 64]) h w 影响 第二维
                # noise torch.Size([1, 16, 32, 32])
                # pixel_latents torch.Size([1, 16, 32, 32])
                # prompt = f"The pair of images highlights first clothing showing second clothing's texture with third clothing's shape, high resolution, 4K, 8K; " \
                #         f"[IMAGE1] Synthesis clothing with second's texture and third's shape." \
                #         f"[IMAGE2] Detailed texture shot of a clothing." \
                #         f"[IMAGE3] Detailed shape shot of a clothing."
    
                model_pred = flux_transformer( # torch.Size([1, 12288, 64])
                    hidden_states=packed_noisy_model_input,
                    timestep=timesteps / 1000,
                    guidance=guidance_vec,
                    pooled_projections=pooled_prompt_embeds,
                    encoder_hidden_states=prompt_embeds,
                    txt_ids=text_ids,
                    # prompt = prompt,
                    img_ids=latent_image_ids,
                    return_dict=False,
                )[0]
                # pdb.set_trace()
                model_pred = FluxFillPipeline._unpack_latents(
                    model_pred,
                    height=noisy_model_input.shape[2] * vae_scale_factor,
                    width=noisy_model_input.shape[3] * vae_scale_factor,
                    vae_scale_factor=vae_scale_factor,
                )
                # these weighting schemes use a uniform timestep sampling
                # and instead post-weight the loss
                weighting = compute_loss_weighting_for_sd3(weighting_scheme=args.weighting_scheme, 
                                                           sigmas=sigmas)

                # flow-matching loss
                target = noise - pixel_latents
                # target = pixel_latents - noise
                loss = torch.mean(
                    (weighting.float() * (model_pred.float() - target.float()) ** 2).reshape(target.shape[0], -1),
                    1,
                )
                loss = loss.mean()
                accelerator.backward(loss)

                if accelerator.sync_gradients:
                    params_to_clip = flux_transformer.parameters()
                    accelerator.clip_grad_norm_(params_to_clip, args.max_grad_norm)
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()

            # Checks if the accelerator has performed an optimization step behind the scenes
            if accelerator.sync_gradients:
                progress_bar.update(1)
                global_step += 1

                # DeepSpeed requires saving weights on every device; saving weights only on the main process would cause issues.
                if accelerator.distributed_type == DistributedType.DEEPSPEED or accelerator.is_main_process:
                    if global_step % args.checkpointing_steps == 0:
                        # _before_ saving state, check if this save would set us over the `checkpoints_total_limit`
                        if args.checkpoints_total_limit is not None:
                            checkpoints = os.listdir(args.output_dir)
                            checkpoints = [d for d in checkpoints if d.startswith("checkpoint")]
                            checkpoints = sorted(checkpoints, key=lambda x: int(x.split("-")[1]))

                            # before we save the new checkpoint, we need to have at _most_ `checkpoints_total_limit - 1` checkpoints
                            if len(checkpoints) >= args.checkpoints_total_limit:
                                num_to_remove = len(checkpoints) - args.checkpoints_total_limit + 1
                                removing_checkpoints = checkpoints[0:num_to_remove]

                                logger.info(
                                    f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
                                )
                                logger.info(f"removing checkpoints: {', '.join(removing_checkpoints)}")

                                for removing_checkpoint in removing_checkpoints:
                                    removing_checkpoint = os.path.join(args.output_dir, removing_checkpoint)
                                    shutil.rmtree(removing_checkpoint)

                        save_path = os.path.join(args.output_dir, f"checkpoint-{global_step}")
                        accelerator.save_state(save_path)
                        logger.info(f"Saved state to {save_path}")

            logs = {"loss": loss.detach().item(), "lr": lr_scheduler.get_last_lr()[0]}
            progress_bar.set_postfix(**logs)
            accelerator.log(logs, step=global_step)

            if global_step >= args.max_train_steps:
                break

    # Create the pipeline using using the trained modules and save it.
    accelerator.wait_for_everyone()
    if accelerator.is_main_process:
        flux_transformer = unwrap_model(flux_transformer)
        if args.upcast_before_saving:
            flux_transformer.to(torch.float32)
        transformer_lora_layers = get_peft_model_state_dict(flux_transformer)
        if args.train_norm_layers:
            transformer_norm_layers = {
                f"transformer.{name}": param
                for name, param in flux_transformer.named_parameters()
                if any(k in name for k in NORM_LAYER_PREFIXES)
            }
            transformer_lora_layers = {**transformer_lora_layers, **transformer_norm_layers}
        FluxFillPipeline.save_lora_weights(
            save_directory=args.output_dir,
            transformer_lora_layers=transformer_lora_layers,
        )

        del flux_transformer
        del vae
        torch.cuda.empty_cache()
        free_memory()

    accelerator.end_training()


if __name__ == "__main__":
    args = parse_args()
    main(args)
